package io.litchi.common.tree.util;

import io.litchi.common.tree.TreeNode;
import lombok.*;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 工具类、
 *
 * @author helang
 * @since 2025-09-29
 */
public class TreeUtil {

    private static final Long DEFAULT_ROOT_VALUE = 0L;

    /**
     * 将扁平的 TreeNode 集合转换为树形结构（森林），返回所有根节点组成的集合
     *
     * @param treeNodeList 扁平节点列表
     * @param rootValue    根节点的父ID（通常是0或null），若为空则使用默认值0
     * @return 构建好的树形结构根节点集合
     */
    public static <T extends TreeNode<T>> Collection<T> toTree(Collection<T> treeNodeList, Long rootValue) {

        if (treeNodeList.isEmpty()) {
            return Collections.emptyList();
        }

        // 如果未传入 rootValue，则使用默认值
        Long actualRootValue = Objects.isNull(rootValue) ? DEFAULT_ROOT_VALUE : rootValue;

        // 使用 Map 存储所有节点，便于快速查找
        Map<Long, T> nodeMap = treeNodeList.stream()
                .collect(Collectors.toMap(TreeNode::getId, node -> node));

        // 先初始化所有节点的 children 集合
        for (T node : treeNodeList) {
            if (node.getChildren() == null) {
                node.setChildren(new ArrayList<>());
            }
        }

        // 构建父子关系
        List<T> roots = new ArrayList<>();
        for (T node : treeNodeList) {
            Long parentId = node.getPid();

            if (Objects.equals(parentId, actualRootValue)) {
                roots.add(node);
            } else if (nodeMap.containsKey(parentId)) {
                T parent = nodeMap.get(parentId);
                parent.getChildren().add(node);
            }
            // 如果 parent 不存在，该节点会被忽略（或可加入孤儿节点处理）
        }

        // 设置层级（level），从根开始 BFS 或 DFS 计算
        for (T root : roots) {
            root.setLevel(0);
            handler(root, 0);
        }

        return roots;
    }

    /**
     * 递归设置每个节点的 level 并可扩展其他逻辑
     *
     * @param node  当前节点
     * @param level 当前层级
     */
    private static <T extends TreeNode<T>> void handler(T node, int level) {
        node.setLevel(level);
        Collection<T> children = node.getChildren();
        if (children != null) {
            for (T child : children) {
                handler(child, level + 1);
            }
        }
    }


    @EqualsAndHashCode(callSuper = true)
    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    public static class  TestTreeNode extends TreeNode{
        private String name;
    }

    public static void main(String[] args) {
        Collection<TestTreeNode> list = new ArrayList<>();
        TestTreeNode testTreeNode = new TestTreeNode();
        testTreeNode.setId(1L);
        testTreeNode.setPid(0L);
        testTreeNode.setName("index1");
        list.add(testTreeNode);
        TestTreeNode testTreeNode2 = new TestTreeNode();
        testTreeNode2.setId(2L);
        testTreeNode2.setPid(0L);
        testTreeNode2.setName("index2");
        list.add(testTreeNode2);

        TestTreeNode testTreeNode3 = new TestTreeNode();
        testTreeNode3.setId(3L);
        testTreeNode3.setPid(1L);
        testTreeNode3.setName("index1-1");
        list.add(testTreeNode3);
        TestTreeNode testTreeNode4 = new TestTreeNode();
        testTreeNode4.setId(4L);
        testTreeNode4.setPid(2L);
        testTreeNode4.setName("index2-1");
        list.add(testTreeNode4);
        TestTreeNode testTreeNode5 = new TestTreeNode();
        testTreeNode5.setId(5L);
        testTreeNode5.setPid(3L);
        testTreeNode5.setName("index1-1-1");
        list.add(testTreeNode5);

        Collection<TestTreeNode> tree = toTree(list, DEFAULT_ROOT_VALUE);
    }
}